import numpy as np
import pandas as pd
import os.path
import os
import time
import tensorflow as tf
# import tensorflow_addons as tfa
from midas import Midas

runs = 1

path = 'speed/data_tmp/'

def waiting(path):
    # dots = "."
  while not os.path.exists(path+'rows_marker.csv'):
    if os.path.exists(path+'rows_marker.csv'):
      break
    time.sleep(0.5)
# print("\rWaiting"+dots)
  print("\rInit...       ")


def reverse_dummies(imputation, cat_cols):
  target_cols = imputation[cat_cols]
  imputation.drop(cat_cols, axis=1, inplace= True)
  imputation['y_1'] = target_cols.idxmax(axis=1)
  return imputation

final_vars = ["birthyr","citylength_1","gender","pew_bornagain","cit1","investor","trans","votereg","pew_religimp","ideo5","pew_churatd","newint","sexuality","educ","internethome","internetwork","marstat","pid3","ownhome","urbancity","immstat","unionhh"]

cont_vars = ["birthyr", "citylength_1"]

while not os.path.exists(path+'rows_marker.csv'):

  if runs > 100:
  	break

  waiting(path)

  data_0 = pd.read_csv(path+"cces_rows_midas_data.csv")

  data_0.columns.str.strip()

  categorical = np.setdiff1d(data_0.columns, cont_vars)

  data_1 = data_0[categorical]
  data_0.drop(categorical, axis= 1, inplace= True)

  constructor_list = [data_0]
  columns_list = []

  for column in data_1.columns:
    na_temp = data_1[column].isnull()
    temp = pd.get_dummies(data_1[column], prefix = column)
    temp[na_temp] = np.nan
    constructor_list.append(temp)
    columns_list.append(list(temp.columns.values))

  data_0 = pd.concat(constructor_list, axis=1)

  na_loc = data_0.isnull()
  data_0[na_loc] = np.nan

  imputer = Midas(layer_structure= [256,256], vae_layer= False, seed= 89, input_drop = 0.75)
  imputer.build_model(data_0, softmax_columns= columns_list)
  imputer.train_model(training_epochs = 20)
  imputer.generate_samples(m=10).output_list
  os.remove(path+'rows_marker.csv')
  
  runs += 1
